from langgraph.prebuilt import create_react_agent
from langchain_ollama import ChatOllama

from langchain_core.messages import convert_to_messages

def pretty_print_message(message, indent=False):
    pretty_message = message.pretty_repr(html=True)
    if not indent:
        print(pretty_message)
        return

    indented = "\n".join("\t" + c for c in pretty_message.split("\n"))
    print(indented)

def pretty_print_messages(update, last_message=False):
    is_subgraph = False
    if isinstance(update, tuple):
        ns, update = update
        # skip parent graph updates in the printouts
        if len(ns) == 0:
            return

        graph_id = ns[-1].split(":")[0]
        print(f"Update from subgraph {graph_id}:")
        print("\n")
        is_subgraph = True

    for node_name, node_update in update.items():
        update_label = f"Update from node {node_name}:"
        if is_subgraph:
            update_label = "\t" + update_label

        print(update_label)
        print("\n")

        messages = convert_to_messages(node_update["messages"])
        if last_message:
            messages = messages[-1:]

        for m in messages:
            pretty_print_message(m, indent=is_subgraph)
        print("\n")

llm_model = ChatOllama(model="qwen2.5:7b-instruct-q5_K_S", temperature=0.2)

# Create a research agent with web search capabilities
def web_search(query: str):
    """Tool that queries GDP data from web and gets back json."""
    search_results = {
        'query': 'What happened at the last wimbledon',
        'follow_up_questions': None,
        'answer': None,
        'images': [],
        'results': 
            [
                {
                    'title': "US GDP and New York State GDP in 2024",
                    'url': "https://www.nbcnews.com/news/sports/andy-murray-wimbledon-tennis-singles-draw-rcna159912",
                    'content': "In 2024, the US GDP was $29.18 trillion and New York State's GDP was $2.297 trillion.",
                    'score': 0.6755297,
                    'raw_content': None,
                }
            ],
        'response_time': 1.31
    }
    return search_results

research_agent = create_react_agent(
    model=llm_model,
    tools=[web_search],
    prompt=(
        "You are a research agent.\n\n"
        "INSTRUCTIONS:\n"
        "- Assist ONLY with research-related tasks, DO NOT do any math\n"
        "- After you're done with your tasks, respond to the supervisor directly\n"
        "- Respond ONLY with the results of your work, do NOT include ANY other text."
    ),
    name="research_agent",
)

# Create a math agent with basic arithmetic operations
def add(a: float, b: float):
    """Add two numbers."""
    return a + b


def multiply(a: float, b: float):
    """Multiply two numbers."""
    return a * b


def divide(a: float, b: float):
    """Divide two numbers."""
    return a / b


math_agent = create_react_agent(
    model=llm_model,
    tools=[add, multiply, divide],
    prompt=(
        "You are a math agent.\n\n"
        "INSTRUCTIONS:\n"
        "- Assist ONLY with math-related tasks\n"
        "- After you're done with your tasks, respond to the supervisor directly\n"
        "- Respond ONLY with the results of your work, do NOT include ANY other text."
    ),
    name="math_agent",
)

# for chunk in math_agent.stream(
#     {"messages": [{"role": "user", "content": "what's (3 + 5) x 7"}]}
# ):
#     pretty_print_messages(chunk)

# Create a supervisor to manage the agents
from langgraph_supervisor import create_supervisor

supervisor = create_supervisor(
    model=llm_model,
    agents=[research_agent, math_agent],
    prompt=(
        "You are a supervisor managing two agents:\n"
        "- a research agent. Assign research-related tasks to this agent\n"
        "- a math agent. Assign math-related tasks to this agent\n"
        "Assign work to one agent at a time, do not call agents in parallel.\n"
        "Do not do any work yourself."
    ),
    add_handoff_back_messages=True,
    output_mode="full_history",
).compile()

for chunk in supervisor.stream(
    {
        "messages": [
            {
                "role": "user",
                "content": "find US and New York state GDP in 2024. what % of US GDP was New York state?",
            }
        ]
    },
):
    pretty_print_messages(chunk, last_message=True)

final_message_history = chunk["supervisor"]["messages"]
print("=================== final_message_history ==========================")
print(final_message_history)
